# 管理不同电脑上的Docker API，并且形成一个独立的FLASK服务，便于其它电脑进行查询。
import logging
import os
import psutil

import docker
import docker.errors
from docker.models.containers import Container
import requests

from backend.models import (
    ContainerResourceUsageDict,
    ContainerInfoDict,
    CreateContainerInfoDict,
    HostMachineResourceInfoDict,
    HostMachineCPUInfo,
    HostMachineMemoryInfo,
    HostMachineDiskInfo,
)
from ..configs import get_config
from typing import Dict, List, Optional, TypedDict, cast, TYPE_CHECKING
from flask import Flask, json, request, jsonify
from .tasks_manager import Task, TaskManager
from .docker_output_format import PullingLogManager
from .utils import get_free_port

logger = logging.getLogger(__name__)


def is_image_exist(image_name: str) -> bool:
    """
    Check if the image is exist in the local docker registry
    """
    try:
        image = client.images.get(image_name)
        return True
    except docker.errors.ImageNotFound or requests.exceptions.HTTPError:
        return False


class CreateContainerTaskResult(TypedDict):
    success: bool
    message: str
    container_info: ContainerInfoDict


# 新增 TypedDict 类型
class RemoveContainerResult(TypedDict):
    success: bool
    message: str


class NetworkInfoDict(TypedDict):
    attachable: bool
    containers: Dict[str, dict]
    driver: str
    id: Optional[str]
    ingress: bool
    internal: bool
    ipam: dict[str, List[dict]]
    name: Optional[str]
    scope: str


class CreateContainerTask(Task):
    result: CreateContainerTaskResult

    def __init__(self, task_id, info: CreateContainerInfoDict):
        super().__init__(task_id, info)
        self.info = info

        self.pulling_log_mgr: PullingLogManager = None

    def get_status_html(self):
        if self.pulling_log_mgr is None:
            return self.log_buffer
        return [self.pulling_log_mgr.render_pulling_progress_html(), *self.log_buffer]

    def run(self):
        if not is_image_exist(self.info["image"]):
            logger.info("Image not found, pulling...")
            self.log_buffer.append("Image not found, pulling...")
            stream = client.api.pull(
                self.info["image"],
                stream=True,
            )
            self.pulling_log_mgr = PullingLogManager(stream)
            self.pulling_log_mgr.process_json_stream()
            self.log_buffer.append(
                "<p>Pulling successful, now creating container...</p>"
            )
        else:
            self.log_buffer.append(
                "<p>Image already exists, now creating container...</p>"
            )

        try:
            container = cast(
                Container,
                client.containers.run(
                    image=self.info["image"],
                    name=self.info["name"],
                    command=self.info["command"],
                    environment=self.info["environment"],
                    ports=cast(dict, self.info["ports"]),
                    network=self.info.get("network", None),
                    networking_config=self.info.get("networking_config", None),
                    detach=True,
                    volumes=[
                        "/var/run/docker.sock:/var/run/docker.sock",
                    ]
                    + self.info.get("volumes", []),
                ),
            )
            self.log_buffer.append(
                msg := f"Successfully created container named '{self.info['name']}' with id '{container.id}', refresh the page to view the container."
            )
            self.result = {
                "success": True,
                "message": msg,
                "container_info": format_container_info(container),
            }
        except docker.errors.APIError as e:
            self.log_buffer.append(msg := f"Error: {e}")
            self.result = {"success": False, "message": msg, "container_info": None}


tm = TaskManager(pool_size=10)


def get_env_variable(container: Container) -> dict[str, str]:
    """
    Get environment variables and convert the original list to a dictionary.
    """
    envs_vars: list[str] = container.attrs["Config"]["Env"]
    return {env.split("=")[0]: env.split("=", maxsplit=1)[1] for env in envs_vars}


def get_image_name(container: Container) -> str:
    """
    Get the image name of the container.

    """
    img = container.image
    if not img:
        return "<None>"
    print(img.tags)
    tag_str = "', '".join(list(set(img.tags)))
    return tag_str


def get_random_free_port_mapping(
    port_mapping: dict[str, int | str],
) -> tuple[dict, dict]:
    """
    If there are values starting with 'FREE_RANDOM_' in the port_mapping dictionary,
    find a random free port for each and return the completed port mapping dictionary
    and a dictionary of the newly assigned random ports.
    """
    random_ports_allocated = {}
    completed_port_mapping = {}
    temp_allocated_ports = set()
    for key, value in port_mapping.items():

        if isinstance(value, str) and value.startswith("FREE_RANDOM_"):
            # Generate a random port number
            random_port = get_free_port(40000, 50000, temp_allocated_ports)
            temp_allocated_ports.add(random_port)
            random_ports_allocated[value] = random_port
            completed_port_mapping[key] = random_port
        else:
            completed_port_mapping[key] = value
        print(key, value, completed_port_mapping)

    return completed_port_mapping, random_ports_allocated


client = docker.from_env()


def list_containers(all=False) -> list[Container]:
    return [container for container in client.containers.list(all=all)]


def start_container(container_name):
    container = client.containers.get(container_name)
    container.start()


def stop_container(container_name):
    container = client.containers.get(container_name)
    container.stop()


def get_host_machine_resource_info() -> HostMachineResourceInfoDict:
    """
    获取当前Linux机器的CPU、内存和磁盘的总量及占用情况。
    """
    # 获取CPU核心数（物理核数和逻辑线程数）
    cpu_physical_count = psutil.cpu_count(logical=False)  # 物理核数
    cpu_logical_count = psutil.cpu_count(logical=True)  # 逻辑线程数

    # 获取每个CPU核心的占用率
    cpu_percentages = psutil.cpu_percent(interval=1, percpu=True)

    # 获取内存信息
    memory = psutil.virtual_memory()
    total_memory = memory.total
    used_memory = memory.used

    # 获取磁盘信息
    disk_partitions = psutil.disk_partitions(True)
    disk_info = []

    for partition in disk_partitions:
        try:
            usage = psutil.disk_usage(partition.mountpoint)
            # Ignore partitions with less than 1GB of total space
            if usage.total < 1024**3:
                continue
            docker_storage_path = get_docker_storage_path()

            # 如果分区的挂载点位于Docker的数据目录下，则忽略该目录
            if (
                os.path.commonprefix([partition.mountpoint, docker_storage_path])
                == docker_storage_path
            ):
                continue
            # 如果路径以/dev等开头，则忽略该目录
            if partition.mountpoint.startswith(
                (
                    "/dev",
                    "/proc",
                    "/sys",
                    "/run",
                    "/snap",
                    get_config()["shared_volume_dir"],
                    "/etc/resolv.conf",
                    "/etc/hosts",
                    "/etc/hostname",
                )
            ):
                continue
            elif partition.mountpoint.endswith(("/docker_instance",)):
                continue

            # 检查Docker数据目录是否在当前分区
            if (
                os.path.commonpath([partition.mountpoint, get_docker_storage_path()])
                == partition.mountpoint
            ):
                is_docker_dir_on_disk = True
            else:
                is_docker_dir_on_disk = False
            disk_info.append(
                {
                    "device": partition.device,
                    "mountpoint": partition.mountpoint,
                    "total": usage.total,
                    "used": usage.used,
                    "free": usage.free,
                    "percent": usage.percent,
                    "is_docker_dir_on_disk": is_docker_dir_on_disk,
                }
            )
        except PermissionError:
            continue

    return {
        "online": True,
        "cpu": cast(
            HostMachineCPUInfo,
            {
                "physical_cores": cpu_physical_count,  # 物理核数
                "logical_cores": cpu_logical_count,  # 逻辑线程数
                "cpu_percentages": cpu_percentages,  # 每个CPU核心的占用率
            },
        ),
        "memory": cast(
            HostMachineMemoryInfo, {"total": total_memory, "used": used_memory}
        ),
        "disks": cast(list[HostMachineDiskInfo], disk_info),
    }


# 修改函数：获取Docker卷的存储路径和Docker文件夹的存储位置
def get_docker_storage_path() -> str:
    """
    Get the storage location of the Docker.
    """
    try:
        daemon_info = client.info()
        docker_root = daemon_info["DockerRootDir"]
    except Exception as e:
        logger.warning(f"Failed to get Docker root directory: {e}")

    return docker_root


def get_containers_status() -> list[ContainerResourceUsageDict]:
    """
    获取所有正在运行的容器的状态信息，包括 CPU、内存、网络、磁盘使用情况。
    """
    containers = client.containers.list()
    result: list[ContainerResourceUsageDict] = []

    for container in containers:
        container_id = container.id
        stats = client.api.stats(container=container_id, stream=False, one_shot=True)
        # print("stats_gen",stats)
        try:
            # 读取一次统计数据
            stats_data = stats  # s json.loads(stats)
            print(stats_data)

            # 提取内存信息
            mem_usage = stats_data["memory_stats"].get("usage", 0)
            mem_limit = stats_data["memory_stats"].get("limit", 0)
            mem_percent = (mem_usage / mem_limit) * 100 if mem_limit > 0 else 0.0

            # 提取 CPU 信息
            cpu_stats = stats_data.get("cpu_stats", {})
            precpu_stats = stats_data.get("precpu_stats", {})

            cpu_delta = cpu_stats.get("cpu_usage", {}).get(
                "total_usage", 0
            ) - precpu_stats.get("cpu_usage", {}).get("total_usage", 0)
            system_delta = cpu_stats.get("system_cpu_usage", 0) - precpu_stats.get(
                "system_cpu_usage", 0
            )
            number_cpus = cpu_stats.get(
                "online_cpus",
                len(cpu_stats.get("cpu_usage", {}).get("percpu_usage", [])),
            )
            cpu_percent = 0.0
            if system_delta > 0 and cpu_delta > 0:
                cpu_percent = (cpu_delta / system_delta) * number_cpus * 100

            # 提取网络信息
            network_rx = (
                stats_data.get("networks", {}).get("eth0", {}).get("rx_bytes", 0)
            )
            network_tx = (
                stats_data.get("networks", {}).get("eth0", {}).get("tx_bytes", 0)
            )

            # 提取磁盘信息
            blkio_stats = stats_data.get("blkio_stats", {}).get(
                "io_service_bytes_recursive", []
            )

            disk_read = (
                sum(
                    item.get("value", 0)
                    for item in blkio_stats
                    if item.get("op") == "Read"
                )
                if blkio_stats
                else 0
            )
            disk_write = (
                sum(
                    item.get("value", 0)
                    for item in blkio_stats
                    if item.get("op") == "Write"
                )
                if blkio_stats
                else 0
            )

            usage_info: ContainerResourceUsageDict = {
                "cpu_usage_percent": round(cpu_percent, 2),
                "mem_usage_bytes": mem_usage,
                "mem_limit_bytes": mem_limit,
                "mem_usage_percent": round(mem_percent, 2),
                "network_rx_bytes": network_rx,
                "network_tx_bytes": network_tx,
                "disk_read_bytes": disk_read,
                "disk_write_bytes": disk_write,
            }

            result.append(usage_info)

        except StopIteration:
            continue

    return result


client = docker.from_env()

app = Flask(__name__)


def format_container_info(container: Container) -> ContainerInfoDict:
    """
    Format container information into a dictionary.
    """
    # Mounts convert to lists:
    mounted_list = [
        f"{mount['Source']}:{mount['Destination']}"
        for mount in container.attrs["Mounts"]
    ]

    # format ports to a dict like
    # original: {'8080/tcp': [{'HostIp': '0.0.0.0', 'HostPort': '40964'}, {'HostIp': '::', 'HostPort': '40964'}]}
    # target: {'8080/tcp': "40964"}
    def get_ports_map(container_ports) -> dict[str, int]:
        ports_map = {}
        for port_key, port_value in container_ports.items():
            if port_value:
                ports_map[port_key] = int(port_value[0]["HostPort"])
        return ports_map

    return {
        "id": cast(str, container.id),
        "name": cast(
            str,
            container.name,
        ),
        "status": container.status,
        "image": get_image_name(container),
        "ports": get_ports_map(container.ports),
        "command": cast(list[str], container.attrs["Config"]["Cmd"]),
        "environment": get_env_variable(container),
        "volumes": mounted_list,
        "network": container.attrs["NetworkSettings"]["Networks"],
    }


@app.route("/containers", methods=["GET"])
def list_containers_api():
    all_containers = request.args.get("all", default=False, type=bool)
    containers = list_containers(all=all_containers)
    return jsonify([format_container_info(container) for container in containers])


@app.route("/containers/<string:container_id_or_name>", methods=["GET"])
def get_container_api(container_id_or_name):
    """
    Get detailed information about a specific container by its name or ID.
    """
    try:
        container = client.containers.get(container_id_or_name)
        return jsonify(format_container_info(container))
    except docker.errors.NotFound:
        return jsonify({"error": f"Container {container_id_or_name} not found"}), 404
    except Exception as e:
        return jsonify({"error": str(e)}), 500


@app.route("/containers/start", methods=["POST"])
def start_container_api():
    data = request.json
    container_name = data.get("name")
    if not container_name:
        return jsonify({"error": "Container name is required"}), 400
    start_container(container_name)
    return jsonify({"message": f"Container {container_name} started"})


@app.route("/containers/stop", methods=["POST"])
def stop_container_api():
    data: dict = request.json
    container_name = data.get("name")
    if not container_name:
        return jsonify({"error": "Container name is required"}), 400
    stop_container(container_name)
    return jsonify({"message": f"Container {container_name} stopped"})


@app.route("/containers/create_container_task_start", methods=["POST"])
def create_container_task_start():
    """
    开始创建新的容器
    注意ports参数，里面如果有FREE_RANDOM_开头的键，就寻找一个随机端口，
    创建成功后返回所有的FIND_RANDOM找到的新端口。
    """
    data: CreateContainerInfoDict = request.json
    required_fields = ["image", "name", "command", "environment", "ports"]

    if not all(field in data for field in required_fields):
        return jsonify({"error": "Missing required fields"}), 400

    ports_mapping, random_ports_mapping = get_random_free_port_mapping(data["ports"])
    processed_data = data.copy()
    processed_data["ports"] = ports_mapping
    logger.info(f"Got ports mapping: {ports_mapping}")
    tm.submit(task := CreateContainerTask("", processed_data))

    return jsonify(
        {
            "message": f"Starting to create {data['name']} from {data['image']}",
            "task_id": task.task_id,
            "random_ports_mapping": random_ports_mapping,
        }
    )


@app.route("/containers/create_container_task_status", methods=["GET"])
def get_create_container_task_status():
    """
    获取指定 task_id 的 CreateContainerTask 状态和 HTML 日志。
    """
    task_id = request.args.get("task_id")
    if not task_id:
        return jsonify({"error": "task_id is required"}), 400

    try:
        status_info = tm.get_create_container_task_status(task_id)
        return jsonify(status_info)
    except KeyError as e:
        return jsonify({"error": str(e)}), 404
    except ValueError as e:
        return jsonify({"error": str(e)}), 400


@app.route("/containers/rename", methods=["POST"])
def rename_container_api():
    """
    Rename a container.
    """
    data = request.json
    container_name = data.get("container_name")
    new_name = data.get("new_name")

    if not container_name or not new_name:
        return jsonify({"error": "Both container_name and new_name are required"}), 400
    try:
        client.containers.get(new_name)
        return (
            jsonify({"error": f"New name for container {new_name} already exists"}),
            400,
        )
    except docker.errors.NotFound:
        pass
    # return jsonify({"error": "aaa"}),400
    try:
        container = client.containers.get(container_name)
        container.rename(new_name)
        return jsonify(
            {
                "message": f"Container {container_name} renamed to {new_name} successfully"
            }
        )
    except docker.errors.NotFound:
        return jsonify({"error": f"Container {container_name} not found"}), 404
    except Exception as e:
        logging.error(str(e))
        return jsonify({"error": str(e)}), 500


def export_container_to_image(container_name: str, image_name: str) -> tuple[str, bool]:
    """
    Export a container to a new image.

    :param container_name: The name of the container to export.
    :param image_name: The name of the new image to create.
    :return: True if the operation was successful, False otherwise.
    """
    try:
        client.images.get(image_name)
        return f"Target image name already exists: {image_name}", False
    except docker.errors.ImageNotFound:
        pass

    try:
        container = client.containers.get(container_name)
        container.commit(repository=image_name)
        return "Success", True
    except docker.errors.NotFound:
        print(msg := f"Container {container_name} not found.")
        return msg, False
    except Exception as e:
        print(msg := f"Error exporting container to image: {str(e)}")
        return msg, False


@app.route("/containers/export", methods=["POST"])
def export_container_api():
    """
    Export a container to a new image.
    """
    data = request.json
    container_name = data.get("container_name")
    image_name = data.get("image_name")

    if not container_name or not image_name:
        return (
            jsonify({"error": "Both container_name and image_name are required"}),
            400,
        )

    msg, success = export_container_to_image(container_name, image_name)

    if success:
        return jsonify(
            {
                "message": f"Container {container_name} exported to image {image_name} successfully"
            }
        )
    else:
        return jsonify({"error": msg}), 400


@app.route("/images/layers", methods=["GET"])
def get_image_layers_api():
    """
    Get the layer information of a specified image.

    :return: A list of layer information.
    """
    image_name = request.args.get("image_name")
    if not image_name:
        return jsonify({"error": "Image name is required"}), 400

    try:
        image = client.images.get(image_name)
        layers = image.attrs["RootFS"]["Layers"]
        return jsonify(layers)
    except docker.errors.ImageNotFound:
        return jsonify({"error": f"Image {image_name} not found"}), 404
    except Exception as e:
        return jsonify({"error": str(e)}), 500


@app.route("/host_machine/resource_info", methods=["GET"])
def get_host_machine_resource_info_api():
    """
    获取当前Linux机器的CPU、内存和磁盘的总量及占用情况。
    """
    logger.info("got_host_machine_resource_info")
    resource_info = get_host_machine_resource_info()
    return jsonify(resource_info)


@app.route("/networks", methods=["GET"])
def list_networks_api():
    """
    获取所有 Docker 网络的信息。
    """
    try:
        networks = client.networks.list()
        network_info_list = []
        for network in networks:
            network_info: NetworkInfoDict = {
                "id": network.id,
                "name": network.name,
                "driver": network.attrs["Driver"],
                "scope": network.attrs["Scope"],
                "ipam": network.attrs["IPAM"],
                "internal": network.attrs["Internal"],
                "attachable": network.attrs["Attachable"],
                # If this is true, the network can be used with the --network flag when creating a container.
                "ingress": network.attrs["Ingress"],
                "containers": network.attrs["Containers"],
            }
            network_info_list.append(network_info)
        return jsonify(network_info_list)
    except Exception as e:
        return jsonify({"error": str(e)}), 500


def remove_container(container_name: str) -> RemoveContainerResult:
    """
    Remove a stopped container.
    If the container is running, return an error.
    """
    try:
        container = client.containers.get(container_name)
        if container.status == "running":
            return {
                "success": False,
                "message": f"Container {container_name} is running and cannot be removed.",
            }
        container.remove()
        return {
            "success": True,
            "message": f"Container {container_name} removed successfully.",
        }
    except docker.errors.NotFound:
        return {"success": False, "message": f"Container {container_name} not found."}
    except Exception as e:
        return {
            "success": False,
            "message": f"Error removing container {container_name}: {str(e)}",
        }


@app.route("/containers/remove", methods=["POST"])
def remove_container_api():
    """
    Remove a stopped container.
    If the container is running, return an error.
    """
    data = request.json
    container_name = data.get("name_or_id")
    if not container_name:
        return jsonify({"error": "Container `name_or_id` is required"}), 400
    result = remove_container(container_name)
    return jsonify(result), 409 if not result["success"] else 200


def run_docker_operator_service(port=58080):
    app.run(host="0.0.0.0", port=port)
